#include "types.h"
#include "kmem.h"
#include "print.h"
#include "lib.h"
#include "regs.h"
#include "csr.h"
#include "cache.h"
#include "mmu.h"

struct vmm_map_desc {
    u64 phy_addr;
    u64 vir_addr;
    u64 size;
    u64 flag;
};

static pagetable_t kpgtbl; //内核页目录
extern char etext[];  // 内核代码段结束地址

// 将OS的VMM标志转换为RISCV MMU标志
// A位和R位要同时设置
// D位和W位要同时设置
static u64 mmu_flag_os2rv(u32 flag)
{
    u64 rv_flag = RV_MMU_GLOBAL;

    if (flag & MMU_FLAG_VALID) {
        rv_flag |= RV_MMU_VALID;
    }

    if (flag & MMU_FLAG_KACCESS) {
        rv_flag |= RV_MMU_ACCESS;
    }

    if (flag & MMU_FLAG_UACCESS) {
        rv_flag |= RV_MMU_USER;
    }

    if (flag & MMU_FLAG_READ) {
        rv_flag |= RV_MMU_READ;
    }

    if (flag & MMU_FLAG_WRITE) {
        rv_flag |= RV_MMU_WRITE | RV_MMU_DIRTY;
    }

    if (flag & MMU_FLAG_EXEC) {
        rv_flag |= RV_MMU_EXEC;
    }

    if (flag & MMU_FLAG_CACHE) {
        rv_flag |= RV_MMU_CACHE | RV_MMU_BUFFER;
    } 
    
    if (flag & MMU_FLAG_DEVICE) {
        rv_flag |= RV_MMU_SO;
    }

    return rv_flag;
}

static inline void mmu_invalidate_tlb_all(void)
{
    asm volatile("sfence.vma zero, zero"); // 失效所有TLB条目
    asm volatile("fence.i"); // 确保指令执行完毕
}

// RISCV SV39 采用三级页表,每级页表项大小8B,一个页表包含512条页表项
// 采用SV39的64位虚拟地址被分为5个部分:
//   39..63 -- 必须为0.
//   30..38 -- 9bit VPN2索引.
//   21..29 -- 9bit VPN1索引.
//   12..20 -- 9bit VPN0索引.
//    0..11 -- 12bit页内偏移.
// 页表项中的RWX三位都是0表示还有下一级页表,否则就是最后一级页表,最后一级的页表项属性参与权限判断
static pte_t *mmu_page_walk(pagetable_t pagetable, u64 va, int alloc)
{
    int level;
    pte_t *pte;

    if(va > MAXVA)
        panic("page_walk");

    for(level = 2; level > 0; level--) {
        // 寻找某级页表项位置
        pte = &pagetable[VPN_GET(level, va)];

        if(*pte & RV_MMU_VALID) {
            // 页表项有效则直接取出下一级页表基址
            pagetable = (pagetable_t)PTE2PA(*pte);
        } else {
            // 页表项无效则可以申请新页内存当作页表
            if(!alloc || (pagetable = os_kalloc()) == 0)
                return (pte_t *)0;
            memset(pagetable, 0, PAGE_SIZE);
            // 设置页表项的有效位,并且RWX为0,表示下一级页表有效
            *pte = PA2PTE(pagetable) | RV_MMU_VALID;
        }

#ifdef KERNEL_MMU_DEBUG
        pr_debug("va %p level %d pte_val %p pte_pa %p.\n", va, level, *pte, PTE2PA(*pte));
#endif
    }

    // 返回三级页表项位置
    return &pagetable[VPN_GET(0, va)];
}

// 建立虚拟地址和物理地址映射关系
// 虚拟地址,物理地址,大小这三个会向下做页对齐处理,使用时最好传入页对齐的数值
int os_mmu_page_map(pagetable_t pagetable, u64 va, u64 pa, u64 size, u32 flag)
{
    u64 rv_flag;
    u64 vstart, vend;
    pte_t *pte;

    // OS标志转换成RV标志
    rv_flag = mmu_flag_os2rv(flag);

    if(size == 0)
        panic("page map size = 0");

    // 虚拟地址和大小都向下做页对齐处理
    vstart = PAGE_ROUNDDOWN(va);
    vend = PAGE_ROUNDDOWN(va + size);

    for (; vstart != vend;) {
        // 根据虚拟地址找到三级页表项位置
        if((pte = mmu_page_walk(pagetable, vstart, 1)) == 0)
            return -1;

        // 检查虚拟地址是否已经映射
        if(*pte & RV_MMU_VALID)
            panic("virtual address already mapped");

        // 设置三级页表项
        *pte = PA2PTE(pa) | rv_flag;

#ifdef KERNEL_MMU_DEBUG
        pr_debug("va %p level 0 pte_val %p pte_pa %p.\n", va, *pte, PTE2PA(*pte));
#endif

        // 处理下一页
        vstart += PAGE_SIZE;
        pa += PAGE_SIZE;
    }

    return 0;
}

static void kvm_map(pagetable_t kpgtbl, u64 va, u64 pa, u64 size, u32 flag)
{
    if(os_mmu_page_map(kpgtbl, va, pa, size, flag) != 0)
        panic("kvmmap fail");

    pr_info("kmap pagetable %p va %p pa %p size %p flag 0x%x.\n", kpgtbl, va, pa, size, flag);
}

// 内核映射表
static struct vmm_map_desc kernel_map_desc[] = {
    {
        KERN_BASE,
        KERN_BASE,
        0xdeaddead,
        VMM_FLAG_KTEXT
    },

    {
        0xdeaddead,
        0xdeaddead,
        0xdeaddead,
        VMM_FLAG_KDATA
    },

    {
        UART0_BASE,
        UART0_BASE,
        4 * 1024,
        VMM_FLAG_DEVICE
    },

    {
        TIMER_BASE,
        TIMER_BASE,
        4 * 1024,
        VMM_FLAG_DEVICE
    },

    {
        PLIC_BASE,
        PLIC_BASE,
        3 * 1024 * 1024,
        VMM_FLAG_DEVICE
    },

    {0 , 0, 0, 0}
};

void os_mmu_init(void)
{
    struct vmm_map_desc *desc = kernel_map_desc;

    // 修正内核代码段信息
    kernel_map_desc[0].size     = (u64)etext - KERN_BASE;

    // 修正内核数据段信息
    kernel_map_desc[1].phy_addr = (u64)etext;
    kernel_map_desc[1].vir_addr = (u64)etext;
    kernel_map_desc[1].size     = KERN_STOP - (u64)etext;

    // 创建内核一级页表
    kpgtbl = (pagetable_t)os_kalloc();
    memset(kpgtbl, 0, PAGE_SIZE);

    // 映射内核地址表
    while (desc->size) {
        kvm_map(kpgtbl, desc->vir_addr, desc->phy_addr, desc->size, desc->flag);
        desc++;
    }

    // 页表改变后需要失效TLB条目
    mmu_invalidate_tlb_all();

    // 设置页目录基址并使能MMU
    csr_write(satp, MAKE_SATP(kpgtbl));
}

// FIXME:
// 此函数调用后会先产生Store/AMO page fault
// 但是处理完成后会接着产生Instruction page fault,原因不明,所以在trap中目前只是panic处理
void mmu_test(void)
{
    volatile u64 *addr = (u64 *)~0;
    *addr = 0;
}

void mmu_kdebug(u64 va)
{
    pagetable_t pagetable = kpgtbl;
    pte_t *pte = mmu_page_walk(pagetable, va, 0);

    printk("pte = %p", pte);
    if (pte) {
        printk("*pte = %p", *pte);
    }
}